# -*- coding: utf-8 -*-

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox
from MainWin import Ui_MainWindow

class LogAnalyzer(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        
        # 初始化变量
        self.all_logs = []  # 存储所有日志行
        self.normal_logs = []  # 存储标记为正常的日志
        self.malice_logs = []  # 存储标记为恶意的日志
        self.current_index = -1  # 当前显示的日志索引
        self.history = []  # 存储历史操作，用于回滚
        # 直接读取malice.dict文件内容到变量，避免每次调用函数都打开文件
        try:
            with open('malice.dict', 'r', encoding='utf-8') as dict_file:
                self.malice_keyword = [line.strip() for line in dict_file.readlines()]
        except FileNotFoundError:
            # 如果文件不存在，初始化一个空列表
            self.malice_keyword = []
        
        # 连接信号和槽
        self.ReadLogfile.clicked.connect(self.read_log_file)
        self.Malice.clicked.connect(self.mark_as_malice)
        self.Normal.clicked.connect(self.mark_as_normal)
        self.PreviousItem.clicked.connect(self.rollback_to_previous)
        self.SaveProject.clicked.connect(self.save_project)
        self.switchMode.clicked.connect(self.switch_mode)
        
    def read_log_file(self):
        """打开文件选择框并读取日志文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择日志文件", "", "文本文件 (*.txt *.log);;所有文件 (*.*)"
        )
        
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as file:
                    self.all_logs = [line.strip() for line in file.readlines()]
                
                # 重置计数器和索引
                self.current_index = 0
                self.normal_logs = []
                self.malice_logs = []
                self.history = []
                
                # 更新统计信息
                self.AllLogNum.setText(str(len(self.all_logs)))
                self.CompleteNum.setText("0")
                self.NormalNum.setText("0")
                self.MaliceNum.setText("0")
                
                # 显示第一行日志
                if self.all_logs:
                    # self.textBrowser.setPlainText(self.all_logs[0])
                    # 使用预先加载的关键字列表判断
                    if any(word in self.all_logs[0] for word in self.malice_keyword):
                        Front_color = "ff0000"
                    else:
                        Front_color = "00ff00"
                    self.textBrowser.setHtml("<!DOCTYPE HTML PUBLIC \"-//W3C//DTD HTML 4.0//EN\" \"http://www.w3.org/TR/REC-html40/strict.dtd\">\n"
"<html><head><meta name=\"qrichtext\" content=\"1\" /><style type=\"text/css\">\n"
"p, li { white-space: pre-wrap; line-height: 1.5; }\n"
"</style></head><body style=\" font-family:\'SimSun\'; font-size:9pt; font-weight:400; font-style:normal;\">\n"
f"<p style=\" margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;\"><span style=\" font-size:12pt; font-weight:600; color:#{Front_color};\">{self.all_logs[0]}</span></p></body></html>")
                    
            except Exception as e:
                self.textBrowser.setPlainText(f"读取文件出错: {str(e)}")
    
    def mark_as_malice(self):
        """将当前行标记为恶意并显示下一行"""
        if 0 <= self.current_index < len(self.all_logs):
            # 保存当前行到恶意列表
            current_log = self.all_logs[self.current_index]
            self.malice_logs.append(current_log)
            
            # 记录操作到历史
            self.history.append(("malice", current_log, self.current_index))
            
            # 更新统计信息
            completed = len(self.normal_logs) + len(self.malice_logs)
            self.CompleteNum.setText(str(completed))
            self.MaliceNum.setText(str(len(self.malice_logs)))
            
            # 显示下一行
            self.current_index += 1
            if self.current_index < len(self.all_logs):
                # self.textBrowser.setPlainText(self.all_logs[self.current_index])
                if any(word in self.all_logs[self.current_index] for word in self.malice_keyword):
                    Front_color = "ff0000"
                else:
                    Front_color = "00ff00"
                self.textBrowser.setHtml("<!DOCTYPE HTML PUBLIC \"-//W3C//DTD HTML 4.0//EN\" \"http://www.w3.org/TR/REC-html40/strict.dtd\">\n"
"<html><head><meta name=\"qrichtext\" content=\"1\" /><style type=\"text/css\">\n"
"p, li { white-space: pre-wrap; line-height: 1.5; }\n"
"</style></head><body style=\" font-family:\'SimSun\'; font-size:9pt; font-weight:400; font-style:normal;\">\n"
f"<p style=\" margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;\"><span style=\" font-size:12pt; font-weight:600; color:#{Front_color};\">{self.all_logs[self.current_index]}</span></p></body></html>")
            else:
                self.textBrowser.setPlainText("所有日志已处理完毕")
    
    def mark_as_normal(self):
        """将当前行标记为正常并显示下一行"""
        if 0 <= self.current_index < len(self.all_logs):
            # 保存当前行到正常列表
            current_log = self.all_logs[self.current_index]
            self.normal_logs.append(current_log)
            
            # 记录操作到历史
            self.history.append(("normal", current_log, self.current_index))
            
            # 更新统计信息
            completed = len(self.normal_logs) + len(self.malice_logs)
            self.CompleteNum.setText(str(completed))
            self.NormalNum.setText(str(len(self.normal_logs)))
            
            # 显示下一行
            self.current_index += 1
            if self.current_index < len(self.all_logs):
                # self.textBrowser.setPlainText(self.all_logs[self.current_index])
                if any(word in self.all_logs[self.current_index] for word in self.malice_keyword):
                    Front_color = "ff0000"
                else:
                    Front_color = "00ff00"
                self.textBrowser.setHtml("<!DOCTYPE HTML PUBLIC \"-//W3C//DTD HTML 4.0//EN\" \"http://www.w3.org/TR/REC-html40/strict.dtd\">\n"
"<html><head><meta name=\"qrichtext\" content=\"1\" /><style type=\"text/css\">\n"
"p, li { white-space: pre-wrap; line-height: 1.5; }\n"
"</style></head><body style=\" font-family:\'SimSun\'; font-size:9pt; font-weight:400; font-style:normal;\">\n"
f"<p style=\" margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;\"><span style=\" font-size:12pt; font-weight:600; color:#{Front_color};\">{self.all_logs[self.current_index]}</span></p></body></html>")
            else:
                self.textBrowser.setPlainText("所有日志已处理完毕")
    
    def rollback_to_previous(self):
        """回滚到上一条处理的日志"""
        if self.history:
            # 获取上一次操作
            action_type, log_content, log_index = self.history.pop()
            
            # 从相应列表中移除
            if action_type == "malice":
                if log_content in self.malice_logs:
                    self.malice_logs.remove(log_content)
            elif action_type == "normal":
                if log_content in self.normal_logs:
                    self.normal_logs.remove(log_content)
            
            # 更新当前索引并显示该行
            self.current_index = log_index
            if 0 <= self.current_index < len(self.all_logs):
                self.textBrowser.setPlainText(self.all_logs[self.current_index])
            
            # 更新统计信息
            completed = len(self.normal_logs) + len(self.malice_logs)
            self.CompleteNum.setText(str(completed))
            self.NormalNum.setText(str(len(self.normal_logs)))
            self.MaliceNum.setText(str(len(self.malice_logs)))

    def save_project(self):
        """保存当前项目"""
        # 将 self.malice_logs 保存到 attack.log.txt
        with open('attack.log.txt', 'a+', encoding='utf-8') as file:
            for log in self.malice_logs:
                file.write(log + '\n')
        # 将 self.normal_logs 保存到 safe.log.txt
        with open('safe.log.txt', 'a+', encoding='utf-8') as file:
            for log in self.normal_logs:
                file.write(log + '\n')

    def switch_mode(self):
        """切换分析模式"""
        # 获取 label_7 内容
        current_mode = self.label_7.text()

        # 切换模式
        if current_mode == "人工标记":
            self.label_7.setText("AI标记")
            # 将 AiMark 按钮设置为 enable
            self.AiMark.setEnabled(True)
        else:
            self.label_7.setText("人工标记")
            self.AiMark.setEnabled(False)
        # 弹窗提示 暂时不支持此功能
        QMessageBox.information(self, "提示", "暂时不支持此功能")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = LogAnalyzer()
    window.show()
    sys.exit(app.exec_())